Skip to content

Isolate key checks in transforms#279

Merged
ValerianRey merged 23 commits intomainfrom
decouple_construction_check_in_transforms
Mar 28, 2025
Merged

Isolate key checks in transforms#279
ValerianRey merged 23 commits intomainfrom
decouple_construction_check_in_transforms

Conversation

@PierreQuinton
Copy link
Copy Markdown
Contributor

@PierreQuinton PierreQuinton commented Mar 28, 2025

  • Add method check_and_get_keys to all transforms
  • Move checks related to keys from init to check_and_get_keys, and make them recursive
  • Remove required_keys and output_keys properties
  • Remove check_keys_are call on the input keys in call and remove check_keys_are from TensorDict
  • Replace _compute with call
  • Change some tests of init into equivalent tests of check_and_get_keys
  • Add new tests for check_and_get_keys
  • Extract the creation of the transform in backward and mlt_backward into their _create_transform functions
  • Rename _make_task_transform to _create_task_transform for uniformity
  • Add test_check_create_transform for these functions
  • Add changelog entry

@PierreQuinton PierreQuinton added package: autojac cc: refactor Conventional commit type for any refactoring, not user-facing, and not typing or perf improvements labels Mar 28, 2025
@ValerianRey ValerianRey changed the title Decouple construction and checking in Transforms Isolate key checks in transforms Mar 28, 2025
@ValerianRey ValerianRey self-requested a review March 28, 2025 11:25
Comment thread src/torchjd/autojac/_transform/_differentiate.py Outdated
Comment thread src/torchjd/autojac/_transform/base.py
Comment thread tests/unit/autojac/_transform/test_base.py
@codecov
Copy link
Copy Markdown

codecov bot commented Mar 28, 2025

Codecov Report

All modified and coverable lines are covered by tests ✅

Files with missing lines Coverage Δ
src/torchjd/autojac/_transform/_differentiate.py 100.00% <100.00%> (ø)
src/torchjd/autojac/_transform/accumulate.py 100.00% <100.00%> (ø)
src/torchjd/autojac/_transform/aggregate.py 100.00% <100.00%> (ø)
src/torchjd/autojac/_transform/base.py 100.00% <100.00%> (ø)
src/torchjd/autojac/_transform/diagonalize.py 100.00% <100.00%> (ø)
src/torchjd/autojac/_transform/init.py 100.00% <100.00%> (ø)
src/torchjd/autojac/_transform/select.py 100.00% <100.00%> (ø)
src/torchjd/autojac/_transform/stack.py 100.00% <100.00%> (ø)
src/torchjd/autojac/_transform/tensor_dict.py 100.00% <ø> (ø)
src/torchjd/autojac/backward.py 100.00% <100.00%> (ø)
... and 1 more
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Comment thread src/torchjd/autojac/_transform/select.py Outdated
Comment thread src/torchjd/autojac/mtl_backward.py
Comment thread src/torchjd/autojac/mtl_backward.py Outdated
Copy link
Copy Markdown
Contributor

@ValerianRey ValerianRey left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM.

In future PRs, we should:

  • Do the same thing for runtime checks of TensorDict: move them from the init to a check method to make them optional.
  • Make typing more strict in Transforms. If instead of having Iterable[Tensor] arguments we have for instance list[Tensor], we can safely defer the call to set(...) in check_and_get_keys, which will save some computations.

Note that the check that we removed from TensorDict (test_keys_are) is actually not necessary to replace by extra unit tests: we test all transforms well enough to ensure that their output keys are what we expect. No need for another PR for this.

@ValerianRey ValerianRey merged commit 1f0afca into main Mar 28, 2025
14 checks passed
@ValerianRey ValerianRey deleted the decouple_construction_check_in_transforms branch March 28, 2025 17:13
@ValerianRey ValerianRey mentioned this pull request May 30, 2025
7 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cc: refactor Conventional commit type for any refactoring, not user-facing, and not typing or perf improvements package: autojac

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants